"""
MCP 客户端
实现 MCP 协议的客户端功能
"""

import asyncio
import json
import logging
from typing import Any, Dict, List, Optional

import aiohttp

from .protocol import MCPRequest, MCPResponse, MCPSerializer
from .security import MCPSecurity

class MCPClient:
    """MCP 客户端
    
    提供 MCP 协议的客户端实现:
    - 连接管理
    - 请求发送
    - 响应处理
    - 重试机制
    """
    
    def __init__(self, config: Dict):
        """初始化客户端
        
        Args:
            config: 客户端配置
                server_url: 服务器地址
                client_id: 客户端标识
                security: 安全配置
                retry: 重试配置
        """
        self.config = config
        self.logger = logging.getLogger(__name__)
        
        # 初始化安全管理器
        self.security = MCPSecurity(config["security"])
        
        # 初始化序列化器
        self.serializer = MCPSerializer()
        
        # 会话信息
        self.session_id = None
        self.token = None
        
        # HTTP 会话
        self.http_session = None
        
        # 重试配置
        self.retry_config = config.get("retry", {
            "max_retries": 3,
            "initial_delay": 1,
            "max_delay": 10,
            "backoff_factor": 2
        })
        
    async def connect(self) -> bool:
        """连接服务器
        
        Returns:
            bool: 是否连接成功
        """
        try:
            # 创建 HTTP 会话
            self.http_session = aiohttp.ClientSession()
            
            # 发送连接请求
            async with self.http_session.post(
                f"{self.config['server_url']}/mcp/connect",
                json={
                    "client_id": self.config["client_id"],
                    "scopes": self.config.get("scopes", [])
                }
            ) as response:
                if response.status == 200:
                    data = await response.json()
                    self.session_id = data["session_id"]
                    self.token = data["token"]
                    self.logger.info("Connected to MCP server")
                    return True
                else:
                    self.logger.error(f"Connection failed: {response.status}")
                    return False
                    
        except Exception as e:
            self.logger.error(f"Connection error: {e}")
            return False
            
    async def disconnect(self) -> bool:
        """断开连接
        
        Returns:
            bool: 是否断开成功
        """
        try:
            if self.session_id:
                # 发送断开请求
                async with self.http_session.post(
                    f"{self.config['server_url']}/mcp/disconnect",
                    json={"session_id": self.session_id}
                ) as response:
                    success = response.status == 200
                    
            # 清理会话
            if self.http_session:
                await self.http_session.close()
                
            self.session_id = None
            self.token = None
            
            self.logger.info("Disconnected from MCP server")
            return True
            
        except Exception as e:
            self.logger.error(f"Disconnection error: {e}")
            return False
            
    async def send_request(self, request: MCPRequest) -> Optional[MCPResponse]:
        """发送请求
        
        Args:
            request: MCP 请求
            
        Returns:
            Optional[MCPResponse]: MCP 响应
        """
        if not self.session_id or not self.token:
            raise RuntimeError("Not connected to server")
            
        # 序列化请求
        request_data = self.serializer.serialize_request(request)
        
        # 添加签名
        signed_request = self.security.sign_request(request_data)
        
        # 重试机制
        retries = 0
        delay = self.retry_config["initial_delay"]
        
        while retries <= self.retry_config["max_retries"]:
            try:
                # 发送请求
                async with self.http_session.post(
                    f"{self.config['server_url']}/mcp/request",
                    json=signed_request,
                    headers={"Authorization": f"Bearer {self.token}"}
                ) as response:
                    if response.status == 200:
                        data = await response.json()
                        return self.serializer.deserialize_response(data)
                    elif response.status == 401:
                        # 令牌过期,尝试重新连接
                        if await self.connect():
                            continue
                        else:
                            return None
                    else:
                        self.logger.error(f"Request failed: {response.status}")
                        
            except Exception as e:
                self.logger.error(f"Request error: {e}")
                
            # 重试
            retries += 1
            if retries <= self.retry_config["max_retries"]:
                await asyncio.sleep(delay)
                delay = min(
                    delay * self.retry_config["backoff_factor"],
                    self.retry_config["max_delay"]
                )
                
        return None
        
    async def query(self, query: str, filters: Dict = None, limit: int = 10) -> Optional[List[Dict]]:
        """发送查询请求
        
        Args:
            query: 查询语句
            filters: 过滤条件
            limit: 结果限制
            
        Returns:
            Optional[List[Dict]]: 查询结果
        """
        request = MCPRequest(
            request_id=f"query_{id(query)}",
            type="query",
            data={
                "query": query,
                "filters": filters or {},
                "limit": limit
            }
        )
        
        response = await self.send_request(request)
        if response and response.status == "success":
            return response.data.get("results")
        return None
        
    async def update(self, entity_id: str, updates: Dict) -> bool:
        """发送更新请求
        
        Args:
            entity_id: 实体ID
            updates: 更新内容
            
        Returns:
            bool: 是否更新成功
        """
        request = MCPRequest(
            request_id=f"update_{entity_id}",
            type="update",
            data={
                "entity_id": entity_id,
                "updates": updates
            }
        )
        
        response = await self.send_request(request)
        if response and response.status == "success":
            return response.data.get("updated", False)
        return False
        
    async def analyze(self, data: List, analysis_type: str, params: Dict = None) -> Optional[Dict]:
        """发送分析请求
        
        Args:
            data: 分析数据
            analysis_type: 分析类型
            params: 分析参数
            
        Returns:
            Optional[Dict]: 分析结果
        """
        request = MCPRequest(
            request_id=f"analysis_{analysis_type}_{id(data)}",
            type="analysis",
            data={
                "data": data,
                "type": analysis_type,
                "params": params or {}
            }
        )
        
        response = await self.send_request(request)
        if response and response.status == "success":
            return response.data.get("results")
        return None
        
    async def stream(self, stream_id: str, action: str, params: Dict = None) -> Optional[Any]:
        """发送流处理请求
        
        Args:
            stream_id: 流ID
            action: 操作类型
            params: 操作参数
            
        Returns:
            Optional[Any]: 处理结果
        """
        request = MCPRequest(
            request_id=f"stream_{stream_id}_{action}",
            type="stream",
            data={
                "stream_id": stream_id,
                "action": action,
                "params": params or {}
            }
        )
        
        response = await self.send_request(request)
        if response and response.status == "success":
            return response.data.get("result")
        return None 